/*
 * Copyright (C) 2021-2022 Red Hat, Inc.
 *
 * Written By: Vadim Rozenfeld <vrozenfe@redhat.com>
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met :
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and / or other materials provided with the distribution.
 * 3. Neither the names of the copyright holders nor the names of their contributors
 *    may be used to endorse or promote products derived from this software
 *    without specific prior written permission.
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS ``AS IS'' AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED.IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
 * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 * SUCH DAMAGE.
 */

#include "pch.h"
#include "Session.h"
#include "pipe.h"
#include "psapi.h"

#define INTERACTIVE_SFX L"_interactive"
#define WINLOGON_SFX    L"_winlogon"
#define PIDS_MAX_COUNT (1024 * 10)

CSession::CSession(ULONG Id) : m_PipeServer_interactive(NULL),
                               m_PipeServer_winlogon(NULL),
                               m_PipeServer_clipboard(NULL)
{
    LOGI("Session id: {}", Id);

    ZeroMemory(&m_SessionInfo, sizeof(m_SessionInfo));
    m_SessionInfo.SessionId = Id;
    ZeroMemory(&m_ProcessInfo_interactive, sizeof(m_ProcessInfo_interactive));
    ZeroMemory(&m_ProcessInfo_winlogon, sizeof(m_ProcessInfo_winlogon));
    ZeroMemory(&m_ProcessInfo_clipboard, sizeof(m_ProcessInfo_clipboard));
}


CSession::~CSession(void)
{
    LOGD("CSession destructor called.");
    Close();
}

bool CSession::Init(void)
{
    LOGI("Init session");

    std::wstring PipeName = PIPE_NAME + std::to_wstring(m_SessionInfo.SessionId) + std::wstring(INTERACTIVE_SFX);
    std::wstring AppName = APP_NAME + std::to_wstring(m_SessionInfo.SessionId) + std::wstring(INTERACTIVE_SFX);
    if (!InitPipeServer(&m_PipeServer_interactive, PipeName.c_str())) {
        return FALSE;
    }
    if (!CreateProcessInSession(AppName, const_cast<LPWSTR>(L"winsta0\\Default"),
                                &m_ProcessInfo_interactive, TRUE)) {
        return FALSE;
    }

    PipeName = PIPE_NAME + std::to_wstring(m_SessionInfo.SessionId) + std::wstring(WINLOGON_SFX);
    AppName = APP_NAME + std::to_wstring(m_SessionInfo.SessionId) + std::wstring(WINLOGON_SFX);
    if (!InitPipeServer(&m_PipeServer_winlogon, PipeName.c_str())) {
        return FALSE;
    }
    if (!CreateProcessInSession(AppName, const_cast<LPWSTR>(L"winsta0\\Winlogon"),
                                &m_ProcessInfo_winlogon, TRUE)) {
        return FALSE;
    }
    return TRUE;
}

bool CSession::RunShareMgr()
{
    LOGI("Run share clipboard programme");

    std::wstring pipeName = VMGTSHARE_PIPE_NAME + std::to_wstring(m_SessionInfo.SessionId);
    std::wstring appName = VMGTSHARE_APP_NAME + std::to_wstring(m_SessionInfo.SessionId);
    if (!InitPipeServer(&m_PipeServer_clipboard, pipeName.c_str())) {
        return FALSE;
    }
    if (!CreateProcessInSession(appName, const_cast<LPWSTR>(L"winsta0\\Default"), &m_ProcessInfo_clipboard, FALSE)) {
        return FALSE;
    }
    return TRUE;
}

bool CSession::InitPipeServer(PipeServer** pipeSrv, LPCWSTR pipeName)
{
    if (!(*pipeSrv)) {
        *pipeSrv = new PipeServer(pipeName);
        return (*pipeSrv)->Init();
    }
    return TRUE;
}

void CSession::Close(void)
{
    LOGI("Close session");

    if (m_PipeServer_interactive) {
        m_PipeServer_interactive->Close();
        delete m_PipeServer_interactive;
        m_PipeServer_interactive = NULL;
    }

    if (m_PipeServer_winlogon) {
        m_PipeServer_winlogon->Close();
        delete m_PipeServer_winlogon;
        m_PipeServer_winlogon = NULL;
    }

    if (m_PipeServer_clipboard) {
        m_PipeServer_clipboard->Close();
        delete m_PipeServer_clipboard;
        m_PipeServer_clipboard = NULL;
    }

    if (m_ProcessInfo_interactive.hProcess)
    {
        CloseHandle(m_ProcessInfo_interactive.hProcess);
        m_ProcessInfo_interactive.hProcess = NULL;
    }

    if (m_ProcessInfo_winlogon.hProcess)
    {
        CloseHandle(m_ProcessInfo_winlogon.hProcess);
        m_ProcessInfo_winlogon.hProcess = NULL;
    }

    if (m_ProcessInfo_clipboard.hProcess)
    {
        CloseHandle(m_ProcessInfo_clipboard.hProcess);
        m_ProcessInfo_clipboard.hProcess = NULL;
    }
}

void CSession::Pause(void)
{
}

void CSession::Resume(void)
{
}

bool CSession::GetSystemTokenForSession(DWORD SessionId, HANDLE *hToken)
{
    DWORD* pids = NULL;
    DWORD cbNeeded = PIDS_MAX_COUNT * sizeof(DWORD);
    DWORD cProcesses = 0;
    DWORD i = 0;
    DWORD dwError = ERROR_SUCCESS;
    HANDLE processHandle = NULL;
    HANDLE TokenLocal = NULL;

    pids = new DWORD[PIDS_MAX_COUNT];
    if (!pids) {
        LOGE("Failed to allocate pids memory.");
        return false;
    }

    if (!EnumProcesses(pids, cbNeeded, &cbNeeded)) {
        dwError = GetLastError();
        LOGE("Failed to call EnumProcesses with error: {}", dwError);
        delete[] pids;
        return false;
    }

    cProcesses = cbNeeded / sizeof(DWORD);
    for (i = 0; i < cProcesses; ++i) {
        do {
            DWORD dwPid = pids[i];
            dwError = ERROR_SUCCESS;

            processHandle = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, true, dwPid);
            if (processHandle == NULL) {
                dwError = GetLastError();
                break;
            }
            WCHAR lpname[MAX_PATH];
            if (!GetProcessImageFileName(processHandle, lpname, MAX_PATH)) {
                dwError = GetLastError();
                break;
            }

            if (wcslen(lpname) < wcslen(L"winlogon.exe")) {
                dwError = ERROR_PATH_NOT_FOUND;
                break;
            }

            if (_wcsicmp(lpname + wcslen(lpname) - wcslen(L"winlogon.exe"), L"winlogon.exe") != 0) {
                dwError = ERROR_PATH_NOT_FOUND;
                break;
            }

            if (!OpenProcessToken(processHandle, TOKEN_DUPLICATE | TOKEN_ASSIGN_PRIMARY | TOKEN_QUERY, &TokenLocal)) {
                dwError = GetLastError();
                break;
            }

            DWORD id = 0;
            DWORD length;
            if (!GetTokenInformation(TokenLocal, TokenSessionId, &id, sizeof(id), &length)) {
                dwError = GetLastError();
                break;
            }

            if (id != SessionId) {
                dwError = ERROR_INVALID_PARAMETER;
                break;
            }
        } while (0);

        if (processHandle) {
            CloseHandle(processHandle);
            processHandle = NULL;
        }

        if (dwError == ERROR_SUCCESS) {
            break;
        } else if(TokenLocal) {
            CloseHandle(TokenLocal);
            TokenLocal = NULL;
        }
    }

    if (TokenLocal) {
        *hToken = TokenLocal;
        return true;
    }
    else {
        return false;
    }
}

bool CSession::CreateProcessInSession(const std::wstring & commandLine, LPWSTR lpDesktop,
                                      PROCESS_INFORMATION *ProcessInfo, bool bSystem)
{
    HANDLE hToken = NULL;
    HANDLE hAdmToken = NULL;
    STARTUPINFO si = { sizeof(si) };
    LPVOID lpvEnv = NULL;
    DWORD dwError = ERROR_SUCCESS;
    wchar_t szUserProfileDir[MAX_PATH];
    DWORD cchUserProfileDir = ARRAYSIZE(szUserProfileDir);

    LOGI(L"Create processes in session.");

    if (ProcessInfo->hProcess && WaitForSingleObject(ProcessInfo->hProcess, 0) == WAIT_TIMEOUT)
    {
        return TRUE;
    }

    do {
        DWORD activeSessionId = WTSGetActiveConsoleSessionId();
        if (activeSessionId != m_SessionInfo.SessionId) {
            LOGE("SessionId: {}, active console sessionId: {}", m_SessionInfo.SessionId, activeSessionId);
            break;
        }

        if (bSystem) {
            if (!GetSystemTokenForSession(m_SessionInfo.SessionId, &hToken))
            {
                LOGE("Failed to call GetSystemTokenForSession");
                dwError = ERROR_INVALID_PARAMETER;
                break;
            }
        } else {
            if (!WTSQueryUserToken(m_SessionInfo.SessionId, &hToken))
            {
                LOGE("Failed to call WTSQueryUserToken with error: {}", GetLastError());
                dwError = ERROR_INVALID_PARAMETER;
                break;
            }
        }

        TOKEN_LINKED_TOKEN admin = {};
        DWORD dw = 0;
        if (GetTokenInformation(hToken, (TOKEN_INFORMATION_CLASS)TokenLinkedToken, &admin, sizeof(TOKEN_LINKED_TOKEN), &dw))
        {
            hAdmToken = admin.LinkedToken;
        }
        else
        {
            DuplicateTokenEx(hToken, MAXIMUM_ALLOWED, NULL, SecurityIdentification, TokenPrimary, &hAdmToken);
        }

        if (!CreateEnvironmentBlock(&lpvEnv, hToken, TRUE))
        {
            dwError = GetLastError();
            LOGE("Failed to call CreateEnvironmentBlock with error: {}", dwError);
            break;
        }

        if (!GetUserProfileDirectory(hToken, szUserProfileDir,
            &cchUserProfileDir))
        {
            dwError = GetLastError();
            LOGE("Failed to call GetUserProfileDirectory with error: {}", dwError);
            break;
        }

        si.lpDesktop = lpDesktop;

        if (!CreateProcessAsUser(hAdmToken, NULL, const_cast<wchar_t *>(commandLine.c_str()), NULL, NULL, FALSE,
            CREATE_UNICODE_ENVIRONMENT | CREATE_NO_WINDOW, lpvEnv, szUserProfileDir, &si, ProcessInfo))
        {
            dwError = GetLastError();
            LOGE("Failed to call CreateProcessAsUser with error: {}", dwError);
            break;
        }
    } while (0);

    if (hAdmToken) {
        CloseHandle(hAdmToken);
        hAdmToken = NULL;
    }

    if (hToken) {
        CloseHandle(hToken);
        hToken = NULL;
    }

    if (lpvEnv) {
        DestroyEnvironmentBlock(lpvEnv);
        lpvEnv = NULL;
    }

    if (ProcessInfo->hThread) {
        CloseHandle(ProcessInfo->hThread);
        ProcessInfo->hThread = NULL;
    }

    if (dwError != ERROR_SUCCESS) {
        SetLastError(dwError);
        return false;
    }

    return true;
}

